#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void constantpad_custom_do(uint32_t blockDim, void *stream, 
                                    uint8_t *x, uint8_t *z, 
                                    int32_t pad_left, int32_t pad_right, 
                                    int32_t pad_top, int32_t pad_bottom,
                                    float pad_value);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void constantpad_custom(GM_ADDR x, GM_ADDR z, 
                                                            int32_t pad_left, int32_t pad_right, 
                                                            int32_t pad_top, int32_t pad_bottom, 
                                                            float pad_value);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;

    // 添加实际的填充参数值
    // 示例值，你可以根据需要修改
    int32_t pad_left = 1;
    int32_t pad_right = 2;  
    int32_t pad_top = 3;
    int32_t pad_bottom = 4;

    float pad_value = 0;

    // 输入形状参数
    int32_t inputN = 8;
    int32_t inputC = 1;
    int32_t inputH = 9;
    int32_t inputW = 2045;

    // 计算输出形状参数
    int32_t outputN = inputN;
    int32_t outputC = inputC;
    int32_t outputH = inputH + pad_top + pad_bottom;
    int32_t outputW = inputW + pad_left + pad_right;

    size_t inputByteSize = inputN * inputC * inputH * inputW * sizeof(float);
    size_t outputByteSize = outputN * outputC * outputH * outputW * sizeof(float);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(constantpad_custom, blockDim, x, z, pad_left, pad_right, pad_top, pad_bottom, pad_value); // 传递参数

    WriteFile("./output/output_z.bin", z, outputByteSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)z);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *xHost, *zHost;
    uint8_t *xDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    constantpad_custom_do(blockDim, stream, xDevice, zDevice, pad_left, pad_right, pad_top, pad_bottom, pad_value);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, outputByteSize);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(zHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}